from nested_gen import *
import numpy
from itertools import izip, chain 

def pack(inArr,phi,nLoops=100):

    # Tests if a cell is unexpected
    def is_unexpected_column(vec,Cold,i,j):
        if (not vec[i]) and Cold[i,j]: return True
        elif vec[i] and not Cold[i,j]: return True
        else: return False
        
    def is_unexpected_row(vec,Cold,i,j):
        if (not vec[j]) and Cold[i,j]: return True
        elif vec[j] and not Cold[i,j]: return True
        else: return False
        
    def c_score(cVals,Cold,UArr,cPos,nRows):
        score = 0.
        for i in xrange(nRows):
            if is_unexpected_column(cVals,Cold,i,cPos): score += UArr[i,cPos]
        return score
        
    def r_score(rVals,Cold,UArr,rPos,mCols):
        score = 0.
        for j in xrange(mCols):
            if is_unexpected_row(rVals,Cold,rPos,j): score += UArr[rPos,j]
        return score
        
    def match_vectors(v1,v2):
        for x in xrange(len(v1)):
            if v1[x] != v2[x]: return False
        return True

    # Define constants
    n,m = inArr.shape

    # Create necessary arrays
    inArr = numpy.array(inArr,dtype=bool)
    coldArr = numpy.array(nest(n,m,phi,0),dtype=bool)
    p = get_parameter(phi,n,m)
    uArr = unexpectedness_array(n,m,p)

    colOrder = numpy.arange(m)
    rowOrder = numpy.arange(n)
    
    prevCols = numpy.zeros(m)
    prevRows = numpy.zeros(m)

    # Begin loop
    for s in xrange(nLoops):
        # Check if we've found the ideal ordering
        if match_vectors(prevCols,colOrder) and match_vectors(prevRows,rowOrder):
            return (rowOrder, colOrder)
        
        # Rearrange the input array to match with the current row and column orderings
        shuffleArr = rearrange_array(inArr,rowOrder,colOrder)
        
        ########## COLUMN SORT #############
        colPositions = numpy.zeros(m)
        colScores = numpy.zeros(m)
        
        # Loop through each column in the current ordering
        for ci in xrange(m):
            colVals = shuffleArr[:,ci]
            bestScore = 1000.
            bestPos = ci
            
            # Loop through the possible positions for the column
            for cPosition in xrange(m):
                cScore = c_score(colVals,coldArr,uArr,cPosition,n)
                
                if cScore < bestScore: # Then update the column's best spot
                    bestScore = cScore
                    bestPos = cPosition
                    
            colPositions[ci] = bestPos
            colScores[ci] = bestScore
            
        colSort = numpy.lexsort((colScores,colPositions))
        prevCols = colOrder
        colOrder = colOrder[colSort]
        
        ############ ROW SORT #############
        # Rearrange the input array to match with the current row and column orderings
        shuffleArr = rearrange_array(inArr,rowOrder,colOrder)
        
        rowPositions = numpy.zeros(m)
        rowScores = numpy.zeros(m)
        
        # Loop through each rowumn in the current ordering
        for ri in xrange(n):
            rowVals = shuffleArr[ri,:]
            bestScore = 1000.
            bestPos = ri
            
            # Loop through the possible positions for the rowumn
            for rPosition in xrange(n):
                rScore = r_score(rowVals,coldArr,uArr,rPosition,m)
                
                if rScore < bestScore: # Then update the row's best spot
                    bestScore = rScore
                    bestPos = rPosition
                    
            rowPositions[ri] = bestPos
            rowScores[ri] = bestScore
            
        rowSort = numpy.lexsort((rowScores,rowPositions))
        prevRows = rowOrder
        rowOrder = rowOrder[rowSort]

    return (rowOrder, colOrder)